###西語短信大類分類模型

import os
import pickle as pkl
import re

import nltk
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F


def tokenize_en(text):
    tokenizer = nltk.tokenize.RegexpTokenizer(r"\w+")
    tokens = tokenizer.tokenize(text)
    word_list = [x for x in tokens if x.isalpha()]
    return word_list


def get_root_dir(path):
    path_list=path.split(os.path.sep)
    index=path_list.index("featurelib")
    return os.path.sep.join(path_list[:index+1])


class Config(object):
    """配置参数"""

    def __init__(self, embedding):
        self.ROOT_DIR = get_root_dir(os.path.abspath("."))
        self.CONF_DIR = os.path.join(
            self.ROOT_DIR, "feature_conf", "sms", "un", "sms_es_cate0_v1"
        )
        self.MODEL_CONF_DIR = os.path.join(self.CONF_DIR, "model_conf")
        self.class_list = [
            x.strip()
            for x in open(
                os.path.join(self.MODEL_CONF_DIR, "class.txt"), encoding="utf-8"
            ).readlines()
        ]  # 类别名单

        self.vocab_path = os.path.join(self.MODEL_CONF_DIR, "vocab.pkl")  # 词表
        self.embedding_pretrained = (
            True if embedding != "random" else None
        )  # 预训练词向量
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )  # 设备

        self.dropout = 0.6  # 随机失活
        self.require_improvement = 5000  # 若超过5000batch效果还没提升，则提前结束训练
        self.num_classes = len(self.class_list)  # 类别数
        self.n_vocab = 0  # 词表大小，在运行时赋值
        self.num_epochs = 5  # epoch数
        self.batch_size = 128  # mini-batch大小
        self.pad_size = 32  # 每句话处理成的长度(短填长切)
        self.learning_rate = 1e-3  # 学习率
        self.embed = 300  # 字向量维度
        self.filter_sizes = (2, 3, 4)  # 卷积核尺寸
        self.num_filters = 256  # 卷积核数量(channels数)


"""Convolutional Neural Networks for Sentence Classification"""


class TextCNN(nn.Module):
    def __init__(self, config):
        super(TextCNN, self).__init__()
        self.embedding = nn.Embedding(
            config.n_vocab, config.embed, padding_idx=config.n_vocab - 1
        )
        self.convs = nn.ModuleList(
            [
                nn.Conv2d(1, config.num_filters, (k, config.embed))
                for k in config.filter_sizes
            ]
        )
        self.dropout = nn.Dropout(config.dropout)
        self.fc = nn.Linear(
            config.num_filters * len(config.filter_sizes), config.num_classes
        )

    def conv_and_pool(self, x, conv):
        x = F.relu(conv(x)).squeeze(3)
        x = F.max_pool1d(x, x.size(2)).squeeze(2)
        return x

    def forward(self, x):
        out = self.embedding(x[0])
        out = out.unsqueeze(1)
        out = torch.cat([self.conv_and_pool(out, conv) for conv in self.convs], 1)
        out = self.dropout(out)
        out = self.fc(out)
        return out


class SmsEsCate0V1ClfModel:
    def __init__(self):
        self.config = Config("random")
        np.random.seed(1)
        torch.manual_seed(1)
        torch.cuda.manual_seed_all(1)
        torch.backends.cudnn.deterministic = True  # 保证每次结果一样
        torch.set_num_threads(8)
        self.UNK = "<UNK>"
        self.PAD = "<PAD>"
        self.stops = eval(
            open(
                os.path.join(self.config.MODEL_CONF_DIR, "spanish_stop_word_list"),
                encoding="utf-8",
            ).read()
        )
        self.tokenizer = lambda x: [
            i.strip() for i in tokenize_en(x) if i not in self.stops
        ]
        self.vocab = pkl.load(open(self.config.vocab_path, "rb"))
        self.config.n_vocab = len(self.vocab)
        self.model = TextCNN(self.config).to(self.config.device)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.load_state_dict(
            torch.load(
                os.path.join(self.config.MODEL_CONF_DIR, "TextCNN.ckpt"),
                map_location=device,
            )
        )
        self.model.eval()

    def load_dataset(self, msg_list, pad_size=32):
        words_line_list = []
        seq_len_list = []
        for msg in msg_list:
            content = msg.strip()
            words_line = []
            token = self.tokenizer(content)
            seq_len = len(token)
            if pad_size:
                if len(token) < pad_size:
                    token.extend([self.PAD] * (pad_size - len(token)))
                else:
                    token = token[:pad_size]
                    seq_len = pad_size
            # word to id
            for word in token:
                words_line.append(self.vocab.get(word, self.vocab.get(self.UNK)))
            words_line_list.append(words_line)
            seq_len_list.append(seq_len)
        return words_line_list, seq_len_list

    def predict(self, text_list):
        if len(text_list) == 0:
            return []
        x, seq_l = self.load_dataset(text_list)
        with torch.no_grad():
            x = torch.LongTensor(x).to(self.config.device)
            seq_l = torch.LongTensor(seq_l).to(self.config.device)
            texts = (x, seq_l)
            outputs = self.model(texts)
            predict_re = (
                torch.max(outputs.data, 1)[1].cpu().numpy().astype(int).tolist()
            )
            probs = F.softmax(outputs).tolist()
        return predict_re


if __name__ == "__main__":
    con = Config("random")
    print(con.MODEL_CONF_DIR)
